package net

import (
	"bytes"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"smart-flow/server/config"
	"smart-flow/server/goal"
	"smart-flow/server/logger"
	"smart-flow/server/utils"
	"strconv"
	"time"

	"go.uber.org/zap"

	"github.com/google/uuid"
)

type AgentContext struct {
	SessionId string         `json:"session_id"`
	AesKey    string         `json:"aes_key"`
	InitTime  int64          `json:"init_time"`
	StartTime int64          `json:"start_time"`
	Uid       string         `json:"uid"`
	Ip        string         `json:"ip"`
	Isp       string         `json:"isp"`
	Status    int            `json:"status"` // 0: 未连接 1: 已连接 2:通道已建立
	Target    goal.TargetCtx `json:"target"`
	Duration  int            `json:"duration"`
	Rate      int            `json:"rate"`
	NeedV6    int            `json:"need_v6"`
	SupportV6 bool           `json:"support_v6"`
	Index     int            `json:"index"`
}

type Agent struct {
	Context       AgentContext
	ManageChan    chan int
	ResultChan    chan string
	SvrPrivateKey *rsa.PrivateKey
	PublicKey     *rsa.PublicKey
}

func (a *Agent) InitAgent(sessionId, uid, ip, isp string, target *goal.TargetCtx, duration, rate, index int, needV6 bool) error {
	a.Context.SessionId = sessionId
	a.Context.AesKey = fmt.Sprintf("%x", uuid.New())[:16]
	a.Context.InitTime = time.Now().Unix()
	a.Context.Uid = uid
	a.Context.Ip = ip
	a.Context.Isp = isp
	a.Context.Status = config.AgentInit
	a.Context.Target = *target
	a.Context.Duration = duration
	a.Context.Rate = rate
	if needV6 {
		a.Context.NeedV6 = 1
	} else {
		a.Context.NeedV6 = 0
	}
	a.Context.SupportV6 = true
	a.Context.Index = index
	a.ManageChan = make(chan int, 10)
	a.ResultChan = make(chan string, 100)
	var err error
	a.SvrPrivateKey, err = utils.GetPrivateKey([]byte(config.AppConfig().Console.PrivateKey))
	return err
}

func CommonAgentLoginReq(dataBuffer []byte, c *ConsoleClient) {
	data := new(CmdAgentLoginReq)
	if err := json.Unmarshal(dataBuffer, data); err != nil {
		logger.Logger.Error("[网络] json反序列化错误", zap.Error(err))
		return
	}
	agent := c.GetSession(data.SessionId)
	if agent == nil {
		logger.Logger.Error("[网络] Agent会话未找到", zap.String("sessionId", data.SessionId))
		return
	}
	var err error
	agent.PublicKey, err = utils.GetPublicKey([]byte(data.ClientPubKey))
	if err != nil {
		logger.Logger.Error("[网络] Agent获取公钥错误", zap.String("sessionId", data.SessionId), zap.Error(err))
		agent.ManageChan <- config.AgentExcept
		return
	}
	encryptDataKey, err := utils.RsaEncrypt([]byte(agent.Context.AesKey), agent.PublicKey)
	if err != nil {
		logger.Logger.Error("[网络] Agent数据rsa加密错误", zap.String("sessionId", data.SessionId), zap.Error(err))
		agent.ManageChan <- config.AgentExcept
		return
	}
	msg := new(CmdAgentLoginRet)
	msg.SessionId = agent.Context.SessionId
	msg.EncryptDataKey = base64.StdEncoding.EncodeToString(encryptDataKey)
	msg.Timestamp = strconv.FormatInt(time.Now().Unix(), 10)
	msg.Uid = agent.Context.Uid
	msgByte, _ := json.Marshal(msg)
	payload := utils.BufferPool.Get().(*bytes.Buffer)
	defer utils.Release(payload)
	if !Pack(uint32(config.CommonAgentLoginRet), payload, msgByte) {
		logger.Logger.Error("[网络] Agent构建报文错误", zap.String("sessionId", data.SessionId))
		agent.ManageChan <- config.AgentExcept
		return
	}
	agent.Context.Status = config.AgentConnected
	agent.ManageChan <- config.AgentConnected
	c.Send(payload)
}

func CommonAgentLogoutReq(dataBuffer []byte, c *ConsoleClient) {
	data := new(CmdAgentLogoutReq)
	if err := json.Unmarshal(dataBuffer, data); err != nil {
		logger.Logger.Error("[网络] json反序列化错误", zap.Error(err))
		return
	}
	agent := c.GetSession(data.SessionId)
	if agent == nil {
		return
	}
	agent.ManageChan <- config.AgentDisconnected
	c.DelSession(data.SessionId)
}

func CommonAgentChannelReq(dataBuffer []byte, c *ConsoleClient) {
	data := new(CmdAgentChannelReq)
	if err := json.Unmarshal(dataBuffer, data); err != nil {
		logger.Logger.Error("[网络] json反序列化错误", zap.Error(err))
		return
	}
	encryptDataKey, err := base64.StdEncoding.DecodeString(data.EncryptDataKey)
	if err != nil {
		logger.Logger.Error("[网络] Agent数据base64解码错误", zap.String("sessionId", data.SessionId), zap.Error(err))
		c.DelSession(data.SessionId)
		return
	}
	agent := c.GetSession(data.SessionId)
	if agent == nil {
		logger.Logger.Error("[网络] Agent会话未找到", zap.String("sessionId", data.SessionId))
		return
	}
	decryptDataKey, err := utils.RsaDecrypt(encryptDataKey, agent.SvrPrivateKey)
	if err != nil {
		logger.Logger.Error("[网络] Agent数据rsa解密错误", zap.String("sessionId", data.SessionId), zap.Error(err))
		agent.ManageChan <- config.AgentExcept
		c.DelSession(data.SessionId)
		return
	}
	if string(decryptDataKey) != agent.Context.AesKey {
		logger.Logger.Error("[网络] Agent认证失败", zap.String("originAesKey", string(decryptDataKey)), zap.String("serverAesKey", agent.Context.AesKey))
		agent.ManageChan <- config.AgentExcept
		c.DelSession(data.SessionId)
		return
	}
	msg := new(CmdAgentChannelRet)
	msg.SessionId = agent.Context.SessionId
	msg.Timestamp = strconv.FormatInt(time.Now().Unix(), 10)
	msg.Uid = agent.Context.Uid
	msgByte, _ := json.Marshal(msg)
	payload := utils.BufferPool.Get().(*bytes.Buffer)
	defer utils.Release(payload)
	if !Pack(uint32(config.CommonAgentChannelRet), payload, msgByte) {
		logger.Logger.Error("[网络] Agent构建报文错误", zap.String("sessionId", data.SessionId))
		agent.ManageChan <- config.AgentExcept
		c.DelSession(agent.Context.SessionId)
		return
	}
	c.Send(payload)
}

func CommonAgentChannelAck(dataBuffer []byte, c *ConsoleClient) {
	data := new(CmdAgentChannelAck)
	if err := json.Unmarshal(dataBuffer, data); err != nil {
		logger.Logger.Error("[网络] json反序列化错误", zap.Error(err))
		return
	}
	agent := c.GetSession(data.SessionId)
	if agent == nil {
		logger.Logger.Error("[网络] Agent会话未找到", zap.String("sessionId", data.SessionId))
		return
	}
	agent.Context.Status = config.AgentReady
	agent.ManageChan <- config.AgentReady
}

func CommonExecCmdRet(dataBuffer []byte, c *ConsoleClient) {
	data := new(CmdAgentExecRet)
	if err := json.Unmarshal(dataBuffer, data); err != nil {
		logger.Logger.Error("[网络] json反序列化错误", zap.Error(err))
		return
	}
	encryptData, err := base64.StdEncoding.DecodeString(data.Result)
	if err != nil {
		logger.Logger.Error("[网络] Agent数据base64解码错误", zap.String("sessionId", data.SessionId), zap.Error(err))
		return
	}
	agent := c.GetSession(data.SessionId)
	if agent == nil {
		// logger.Logger.Warn("[网络] Agent会话未找到", zap.String("sessionId", data.SessionId))
		return
	}
	decryptData, err := utils.AesDecrypt(encryptData, []byte(agent.Context.AesKey))
	if err != nil {
		logger.Logger.Error("[网络] Agent数据aes解密错误", zap.String("sessionId", data.SessionId), zap.Error(err))
		agent.ManageChan <- config.AgentExcept
		return
	}
	// logger.Logger.Info("[网络] Agent指令返回", zap.String("cmdRet", string(decryptData)))
	agent.ResultChan <- string(decryptData)
}

func CommonExecCmdReq(sessionId, uid, cmd string, index int) {
	agent := ConsoleInstance(index).GetSession(sessionId)
	if agent == nil {
		logger.Logger.Error("[网络] Agent会话未找到", zap.String("sessionId", sessionId))
		return
	}
	encryptCmd, err := utils.AesEncrypt([]byte(cmd), []byte(agent.Context.AesKey))
	if err != nil {
		logger.Logger.Error("[网络] Agent数据aes加密错误", zap.String("sessionId", sessionId), zap.Error(err))
		agent.ManageChan <- config.AgentExcept
		return
	}
	msg := new(CmdAgentExecReq)
	msg.SessionId = sessionId
	msg.Uid = uid
	msg.Cmd = base64.StdEncoding.EncodeToString(encryptCmd)
	msgByte, _ := json.Marshal(msg)
	payload := utils.BufferPool.Get().(*bytes.Buffer)
	defer utils.Release(payload)
	if !Pack(uint32(config.CommonExecCmdReq), payload, msgByte) {
		logger.Logger.Error("[网络] Agent构建报文错误", zap.String("sessionId", sessionId))
		agent.ManageChan <- config.AgentExcept
		return
	}
	ConsoleInstance(index).Send(payload)
}
